import subprocess
import argparse

def check_iptables_port(port, protocol="tcp", chain="INPUT"):
    """
    检查 Linux iptables 是否放行指定端口
    
    :param port: 端口号 (int)
    :param protocol: 协议类型 ("tcp" 或 "udp") (str)
    :param chain: 链名称 ("INPUT" / "OUTPUT" / "FORWARD") (str)
    :return: True (开放) / False (未开放) / None (错误)
    """
    if protocol.lower() not in ["tcp", "udp"]:
        print("错误：protocol 必须是 'tcp' 或 'udp'")
        return None
    
    if chain.upper() not in ["INPUT", "OUTPUT", "FORWARD"]:
        print("错误：chain 必须是 'INPUT', 'OUTPUT' 或 'FORWARD'")
        return None
    
    try:
        # 使用 subprocess.run，并设置 universal_newlines=True 替代 text=True
        result = subprocess.run(
            ["iptables", "-L", chain, "-n", "--line-numbers"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True  # 替代 text=True
        )
        
        if result.returncode != 0:
            print(f"执行 iptables 命令失败: {result.stderr.strip()}")
            return None
        
        output = result.stdout
        
        # 解析 iptables 输出
        in_accept_section = False
        for line in output.splitlines():
            # 检查是否进入 ACCEPT 规则
            if line.strip().startswith("ACCEPT"):
                in_accept_section = True
                continue
            
            # 如果不在 ACCEPT 规则部分，跳过
            if not in_accept_section:
                continue
            
            # 检查是否是规则行（忽略统计信息和空行）
            if line.strip() == "" or "pkts" in line or "bytes" in line:
                continue
            
            # 检查协议和端口
            parts = line.split()
            if len(parts) < 4:
                continue
            
            # 提取协议
            current_protocol = parts[0].lower()
            if current_protocol != protocol:
                continue
            
            # 提取目标端口
            # iptables 规则中的端口通常在 --dport 参数后
            dport_index = None
            for i, part in enumerate(parts):
                if part == "--dport":
                    dport_index = i + 1
                    break
            
            if dport_index is not None and dport_index < len(parts):
                try:
                    rule_port = int(parts[dport_index])
                    if rule_port == port:
                        print(f"[{chain}] {protocol.upper()} 端口 {port} 已开放")
                        return True
                except ValueError:
                    # 如果端口不是整数，跳过
                    continue
        
        print(f"[{chain}] {protocol.upper()} 端口 {port} 未开放")
        return False
    
    except Exception as e:
        print(f"发生错误: {e}")
        return None

def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description="检查 Linux iptables 是否放行指定端口")
    parser.add_argument("--port", type=int, required=True, help="要检查的端口号（如 80）")
    parser.add_argument("--protocol", type=str, default="tcp", choices=["tcp", "udp"], help="协议类型（tcp/udp，默认 tcp）")
    parser.add_argument("--chain", type=str, default="INPUT", choices=["INPUT", "OUTPUT", "FORWARD"], help="链名称（INPUT/OUTPUT/FORWARD，默认 INPUT）")
    
    args = parser.parse_args()
    
    # 检查端口
    is_open = check_iptables_port(args.port, args.protocol, args.chain)
    
    if is_open is not None:
        if is_open:
            print(f"✅ 端口 {args.port} ({args.protocol.upper()}) 已开放")
        else:
            print(f"❌ 端口 {args.port} ({args.protocol.upper()}) 未开放")

# python3 iptables-check.py  --port 28080 --protocol tcp --chain INPUT
if __name__ == "__main__":
    main()